#!/usr/bin/env python3
# -*- encoding:utf-8 -*-

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
import zlib
import struct
import zipfile
import os
import argparse
import logging


class SsfAES:
    _key = b'\x52\x36\x46\x1A\xD3\x85\x03\x66' + \
        b'\x90\x45\x16\x28\x79\x03\x36\x23' + \
        b'\xDD\xBE\x6F\x03\xFF\x04\xE3\xCA' + \
        b'\xD5\x7F\xFC\xA3\x50\xE4\x9E\xD9'

    _iv = b'\xE0\x7A\xAD\x35\xE0\x90\xAA\x03' + \
        b'\x8A\x51\xFD\x05\xDF\x8C\x5D\x0F'

    @classmethod
    def key(self):
        return self._key

    @classmethod
    def iv(self):
        return self._iv


def extractSsf(ssf_file_path, dest_dir):
    """
        提取 ssf 文件到指定文件夹，文件夹不存在会自动创建
        ssf 文件格式目前有两种，一种是加密过后，一种未加密的zip
        -> bool
    """
    # 保存目录名不能是文件
    if os.path.isfile(dest_dir):
        logging.error(f'用于保存的文件夹 {dest_dir} 不可以是一个已经存在的文件')
        return False

    # 读取文件的二进制内容
    with open(ssf_file_path, 'rb') as ssf_file:
        ssf_bin = ssf_file.read()

    # 文件头4字节判断是否被加密
    if ssf_bin[:4] == b'Skin':

        # AES 解密内容
        ssf_aes = AES.new(SsfAES.key(), AES.MODE_CBC, SsfAES.iv())
        decrypted_ssf_bin = ssf_aes.decrypt(ssf_bin[8:])

        # zlib 解压内容
        data = zlib.decompress(decrypted_ssf_bin[4:])  # 注意要跳过头4字节

        def readUint(offset):
            return struct.unpack('I', data[offset:offset + 4])[0]

        # 整个内容的大小 参考用 占4字节
        # size = readUint(0)

        # 得到每个内容文件偏移量
        offsets_size = readUint(4)
        offsets = struct.unpack('I'*(offsets_size//4),
                                data[8:8 + offsets_size])

        # 创建文件夹
        if not os.path.isdir(dest_dir):
            os.mkdir(dest_dir)

        for offset in offsets:
            # 得到文件名
            name_len = readUint(offset)
            filename = data[offset+4:offset+4+name_len].decode('utf-16')

            # 得到文件内容
            content_len = readUint(offset+4+name_len)
            content = data[offset + 8 +
                           name_len:offset + 8 + name_len+content_len]

            # 写入文件
            with open(dest_dir.rstrip(os.sep) + os.sep + filename, 'wb') as origin_file:
                origin_file.write(content)
        return True

    elif ssf_bin[:2] == b'PK':
        # 解压 zip
        with zipfile.ZipFile(ssf_file_path) as zf:
            zf.extractall(dest_dir)
            return True

    logging.error(f'{ssf_file_path} 文件格式不支持')
    return False


def packEncryptedSsf(dest_dir, ssf_file_path):
    """
        打包目录为加密的ssf文件
        -> bool
    """
    if not os.path.isdir(dest_dir):
        logging.error(f'要打包的文件夹 {dest_dir} 路径不存在或不是文件夹')
        return False

    if os.path.isdir(ssf_file_path) or len(ssf_file_path.strip()) < 1:
        logging.error(f'打包文件夹到 {ssf_file_path} 文件路径必须填写正确并且不能是已存在的文件夹路径')
        return False

    # 整数转 4 字节
    def packIntUint(int_num):
        return struct.pack('I', int_num)

    # 数据长度转 4 字节
    def packDataUint(data):
        return struct.pack('I', len(data))

    # 内容文件偏移转元祖
    def packOffsets(offsets):
        return struct.pack(f'{len(offsets)}I', *offsets)
    ssfData = bytes()

    # 编译 ssf 文件头 (8 bytes)
    ssfHead = 'Skin'.encode('utf-8') + packIntUint(3)

    # 遍历目录文件数据封装到数据部分
    offsets = []
    offset = 0
    for item in os.scandir(dest_dir):
        if item.is_file():
            # 编译文件名
            filename = item.name.encode('utf-16-le')
            name_len = packDataUint(filename)

            # 编译文件内容
            with open(item.path, 'rb') as origin_file:
                content = origin_file.read()
                content_len = packDataUint(content)

            # 合并数据
            filedata = name_len + filename + content_len + content
            ssfData += filedata

            # 记录文件数量和文件相对偏移
            offset += len(filedata)
            offsets.append(offset)

    offsets_size = len(offsets)

    # 修正偏移组
    if offsets_size > 0:
        offsets.insert(0, 0)
        offsets.pop()

    # 更新文件偏移组到绝对偏移
    offsets_head_size = 8 + offsets_size * 4
    for offset in range(offsets_size):
        offsets[offset] += offsets_head_size

    # 合并内容数据
    ssfData = packIntUint(offsets_head_size + len(ssfData)) + \
        packIntUint(offsets_size * 4) + \
        packOffsets(offsets) + ssfData

    # zlib 压缩内容数据
    ssfData = ssfData[:4] + zlib.compress(ssfData)

    # 加密打包数据
    ssf_aes = AES.new(SsfAES.key(), AES.MODE_CBC, SsfAES.iv())
    ssfData = ssf_aes.encrypt(pad(ssfData, 16))
    ssfData = ssfHead + ssfData

    # 把打包数据写入 ssf 文件
    with open(ssf_file_path, 'wb') as ssf_file:
        ssf_file.write(ssfData)
    return True


def main(args):
    logging.basicConfig(
        format='\n%(levelname)s：%(message)s', level=logging.INFO)

    # 如果是文件就提取 否则 如果是文件夹就打包
    if os.path.isfile(args.src):
        logging.info(f'提取文件：{args.src}')
        result = extractSsf(args.src, args.dest)
        if not result:
            logging.error('提取失败！')
            return 1
        logging.info('提取成功！')

    elif os.path.isdir(args.src):
        logging.info(f'打包文件夹：{args.src}')
        result = packEncryptedSsf(args.src, args.dest)
        if not result:
            logging.error('打包失败！')
            return 1
        logging.info('打包成功！')

    else:
        logging.error(f'找不到文件或文件夹：{args.src}')
        return 1

    return 0


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawTextHelpFormatter,
        description='''搜狗输入法皮肤文件 (.ssf) 提取与打包工具
支持搜狗输入法皮肤编辑器版本到：6.6''')

    parser.add_argument('src', help='''
提取：输入 .ssf 皮肤文件路径
打包：输入要打包的皮肤文件夹路径''')

    parser.add_argument('dest', help='''
提取：输入要保存提取皮肤文件的文件夹名称
打包：输入 .ssf 皮肤文件路径''')

    args = parser.parse_args()

    exit(main(args))
